import numpy as np

from implement.layers.conv.deconv2d import Deconv2d
from implement.layers.basic.model import Model
from implement.layers.conv.conv2d import Conv2d
from implement.layers.basic.affine import Affine
from utils.functions_collect import reshape, relu, sigmoid


class Decoder(Model):
    def __init__(self):
        """解码器模型，用于将输入数据进行反向处理生成图像。

        包含一个仿射变换层、一个反卷积层和一个卷积层。

        - 仿射变换层：将输入数据展平成一维数组。
        - 反卷积层：将展平的数据重塑为三维数据。
        - 卷积层：处理反卷积后的数据并输出图像。

        注意：在输出前应用了 ReLU 和 Sigmoid 激活函数。

        Attributes:
            to_shape (tuple): 目标形状，即展平前的形状 (C, H, W)。
            linear (Affine): 仿射变换层。
            deconv (Deconv2d): 反卷积层。
            conv (Conv2d): 卷积层。
        """
        super().__init__()
        self.to_shape = (64, 14, 14)  # (C, H, W)
        self.linear = Affine(np.prod(self.to_shape))
        self.deconv = Deconv2d(32, kernel_size=4, stride=2, pad=1)
        self.conv = Conv2d(1, kernel_size=3, stride=1, pad=1)

    def forward(self, x):
        """解码器的前向传播。

        Args:
            x (ndarray): 输入数据。

        Returns:
            ndarray: 解码器的输出，即生成的图像。
        """
        x = relu(self.linear(x))
        x = reshape(x, (-1,) + self.to_shape)  # reshape to (-1, C, H, W)
        x = relu(self.deconv(x))
        x = self.conv(x)
        x = sigmoid(x)
        return x


